from typing import IO, List, Dict, Any, Optional

from llama_index.core.schema import Document

from axiom_boot.di.decorators import service
from .models import FilePayload
from .data_parser_interface import DataParserInterface
# 导入所有具体的解析器
from .parsers import (
    PdfParser,
    MarkdownParser,
    DocxParser,
    HtmlParser,
    ImageParser,
    UnstructuredParser,
)


@service()
class DocumentParserManager:
    """
    一个通用的、批量优先的内容解析引擎。
    负责将多种形态的原始内容，高效地转换为标准的 LlamaIndex Document 列表。
    """

    def __init__(self):
        """
        初始化管理器，并建立文件格式与解析器之间的映射关系（注册表）。
        """
        # _file_parsers 是一个字典，键是文件格式（小写），值是对应的解析器 *类*。
        # 这使得添加新的文件类型支持变得非常简单，只需在此处添加一行即可。
        self._file_parsers: Dict[str, type[DataParserInterface]] = {
            "pdf": PdfParser,
            "md": MarkdownParser,
            "docx": DocxParser,
            "html": HtmlParser,
            "htm": HtmlParser,  # 同时支持 .htm 后缀
            "jpg": ImageParser,
            "jpeg": ImageParser, # 同时支持 .jpeg 后缀
            "png": ImageParser,
            "txt": UnstructuredParser, 
            # 将 UnstructuredParser 作为处理其他所有未知类型的后备解析器
            "default": UnstructuredParser,
        }

    # --- 场景1: 程序内部生成的文本 ---
    def from_texts(self, texts: List[str], metadatas: Optional[List[Dict[str, Any]]] = None) -> List[Document]:
        """
        [批量] 从内存中的文本字符串列表创建 Documents。
        这是处理程序内部动态生成文本的最高效方式。
        """
        metadatas = metadatas or ([{}] * len(texts))
        if len(texts) != len(metadatas):
            raise ValueError("文本列表 (texts) 的数量必须与元数据列表 (metadatas) 的数量相匹配。")

        return [Document(text=text, metadata=meta) for text, meta in zip(texts, metadatas)]

    # --- 场景2: 各类文件 (PDF, Office, Markdown, HTML, 图片) ---
    def from_files(self, files: List[FilePayload], **kwargs) -> List[Document]:
        """
        [批量] 从文件负载列表进行解析。
        这是处理所有文件类数据的统一入口，无论是 PDF、Word 还是图片 OCR。
        """
        all_documents = []
        # 遍历每一个待处理的文件负载
        for file_payload in files:
            # 获取文件格式并转换为小写，以便在注册表中查找
            file_format = file_payload.file_format.lower()
            
            # 关键修正：不再使用共享的单例解析器，而是在每次循环中创建全新的实例。
            # 这可以从根本上避免因解析器内部状态残留而导致的后续文件处理失败。
            parser_class = self._file_parsers.get(file_format, self._file_parsers["default"])
            parser = parser_class()

            # 调用解析器的 parse 方法来处理文件流
            try:
                # 核心修正：将 file_format 传递给 parse 方法
                documents = parser.parse(
                    file=file_payload.file_stream,
                    file_format=file_format
                )
                
                # 最终修正：我们不仅要更新 metadata 字典，还必须设置 doc_id 属性。
                # LlamaIndex 的 NodeParser 在切分时，会将 doc.doc_id 传递给 node.ref_doc_id，
                # 而 PGVectorStore 在进行元数据过滤时，依赖的正是这个 ref_doc_id。
                if file_payload.metadata:
                    doc_id = file_payload.metadata.get("document_id")
                    for doc in documents:
                        if doc_id:
                            # 关键修正：同时更新 metadata 和 doc_id 属性
                            doc.metadata["document_id"] = doc_id
                            doc.doc_id = doc_id
                        
                all_documents.extend(documents)
            except Exception as e:
                # 在生产环境中，这里应该使用更完善的日志记录和错误处理机制
                print(f"处理文件失败 (格式: {file_format}): {e}")

        return all_documents

    # --- 场景3: 结构化数据 (例如，来自数据库或 CSV) ---
    def from_dicts(
        self,
        records: List[Dict[str, Any]],
        text_template: str
    ) -> List[Document]:
        """
        [批量] 从字典列表创建 Documents。
        这使得将任何结构化数据源集成到知识库中变得异常简单。

        Args:
            records (List[Dict[str, Any]]): 字典列表，每个字典代表一条记录。
            text_template (str): 一个格式化字符串，用于将字典转换为一段有意义的文本。
                                 例如: "用户 {name} (ID: {user_id}) 发表了评论: {comment}"
        """
        documents = []
        # 遍历每一条字典记录
        for record in records:
            # 使用用户提供的模板，通过字符串格式化功能，从记录中动态生成文本内容
            text = text_template.format(**record)
            
            # 我们将原始的、完整的字典记录作为元数据 (metadata) 存储起来。
            # 这样做的好处是，在检索到相关文档后，可以方便地回溯其所有原始信息。
            # 使用 .copy() 是一个好习惯，可以避免后续操作意外修改原始数据。
            metadata = record.copy()
            
            documents.append(Document(text=text, metadata=metadata))
            
        return documents

    # --- 场景4: 多媒体 (音频, 视频) ---
    def from_multimedia(self, files: List[FilePayload], **kwargs) -> List[Document]:
        """
        [批量] 从多媒体文件进行解析，例如音频和视频。
        这展示了框架的未来可扩展性。
        """
        # ... 未来可以实现具体的 ASR (语音转文本) 或视频分析模块 ...
        raise NotImplementedError("多媒体文件解析功能尚未实现。")

    # --- 为方便使用，可以为每个批量方法提供一个单数形式的包装器 ---
    def from_text(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> List[Document]:
        """[便捷方法] 处理单个文本字符串。"""
        return self.from_texts([text], [metadata] if metadata else [{}])

    def from_file(self, file_format: str, file_stream: IO[bytes], **kwargs) -> List[Document]:
        """[便捷方法] 处理单个文件流。"""
        # 修正：同样移除 kwargs 的传递
        payload = FilePayload(file_format=file_format, file_stream=file_stream, metadata=kwargs.get("metadata"))
        return self.from_files([payload]) 